package com.jf.ollama.service;

import com.fasterxml.jackson.annotation.JsonProperty;
import com.fasterxml.jackson.core.JsonProcessingException;
import com.fasterxml.jackson.databind.ObjectMapper;
import com.jf.ollama.bean.MathReasoning;
import org.springframework.ai.chat.model.ChatModel;
import org.springframework.ai.chat.model.ChatResponse;
import org.springframework.ai.chat.prompt.Prompt;
import org.springframework.ai.converter.BeanOutputConverter;
import org.springframework.ai.ollama.OllamaChatModel;
import org.springframework.ai.ollama.api.OllamaModel;
import org.springframework.ai.ollama.api.OllamaOptions;
import org.springframework.stereotype.Service;

import java.util.Map;

/**
 * Created with IntelliJ IDEA
 * Description:
 * User:22954
 * Date:2025-06-15
 * Time:13:12
 */
@Service
public class OutputStructService {
    private final OllamaChatModel chatModel;
    public OutputStructService(OllamaChatModel chatModel) {
        this.chatModel = chatModel;
    }

    public ChatResponse outputStruct(String message) {
        String jsonSchema = """
        {
            "type": "object",
            "properties": {
                "steps": {
                    "type": "array",
                    "items": {
                        "type": "object",
                        "properties": {
                            "explanation": { "type": "string" },
                            "output": { "type": "string" }
                        },
                        "required": ["explanation", "output"],
                        "additionalProperties": false
                    }
                },
                "final_answer": { "type": "string" }
            },
            "required": ["steps", "final_answer"],
            "additionalProperties": false
        }
        """;

        Prompt prompt = null;
        try {
            prompt = new Prompt(message,
                    OllamaOptions.builder()
                            .format(new ObjectMapper().readValue(jsonSchema, Map.class))
                            .build());
        } catch (JsonProcessingException e) {
            throw new RuntimeException(e);
        }

        return chatModel.call(prompt);
    }

    public MathReasoning beanOutputStruct(String message) {
        var outputConverter = new BeanOutputConverter<>(MathReasoning.class);

        Prompt prompt = new Prompt(message,
                OllamaOptions.builder()
                        .format(outputConverter.getJsonSchemaMap())
                        .build());

        ChatResponse response = chatModel.call(prompt);
        String content = response.getResult().getOutput().getText();

        return outputConverter.convert(content);
    }
}
